use anyhow::{anyhow, Result};
use std::{collections::HashMap, net::SocketAddr, time::Duration};
use tokio::sync::{Mutex, RwLock};
use tokio::time::timeout;
use tokio_modbus::{client::tcp, client::Context, prelude::*, slave::Slave};

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct SlaveKey {
    pub addr: SocketAddr,
    pub unit_id: u8,
}

impl SlaveKey {
    pub fn new(addr: SocketAddr, unit_id: u8) -> Self {
        Self { addr, unit_id }
    }
}

#[derive(Clone)]
pub struct ModbusConnectionPool {
    inner: std::sync::Arc<RwLock<HashMap<SlaveKey, std::sync::Arc<Mutex<Option<Context>>>>>>,
    connect_timeout: Duration,
}

impl ModbusConnectionPool {
    pub fn new() -> Self {
        Self {
            inner: std::sync::Arc::new(RwLock::new(HashMap::new())),
            connect_timeout: Duration::from_secs(3),
        }
    }

    pub fn with_connect_timeout(mut self, dur: Duration) -> Self {
        self.connect_timeout = dur;
        self
    }

    async fn get_handle(&self, key: &SlaveKey) -> std::sync::Arc<Mutex<Option<Context>>> {
        // fast path
        {
            let map = self.inner.read().await;
            if let Some(h) = map.get(key) {
                return h.clone();
            }
        }
        let mut map = self.inner.write().await;
        map.entry(key.clone())
            .or_insert_with(|| std::sync::Arc::new(Mutex::new(None)))
            .clone()
    }

    async fn connect(&self, key: &SlaveKey) -> Result<Context> {
        let fut = tcp::connect_slave(key.addr, Slave(key.unit_id));
        match timeout(self.connect_timeout, fut).await {
            Err(_) => Err(anyhow!("connect timeout to {}", key.addr)),
            Ok(Err(e)) => Err(anyhow!("connect error to {}: {}", key.addr, e)),
            Ok(Ok(ctx)) => Ok(ctx),
        }
    }

    // Helper: ensure connected and return a locked context guard
    async fn ensure_ctx(&self, key: &SlaveKey) -> Result<std::sync::Arc<Mutex<Option<Context>>>> {
        let handle = self.get_handle(key).await;
        let mut guard = handle.lock().await;
        if guard.is_none() {
            let ctx = self.connect(key).await?;
            *guard = Some(ctx);
        }
        drop(guard);
        Ok(handle)
    }

    pub async fn read_holding(&self, key: SlaveKey, addr: u16, cnt: u16) -> Result<Vec<u16>> {
        let k = key.clone();
        self.ensure_ctx(&k).await?;
        let handle = self.get_handle(&k).await;
        let mut guard = handle.lock().await;
        let r = match guard.as_mut() {
            Some(ctx) => ctx.read_holding_registers(addr, cnt).await.map_err(|e| anyhow!(e)),
            None => Err(anyhow!("No context")),
        };
        if r.is_err() {
            *guard = None;
        }
        r
    }

    pub async fn read_input(&self, key: SlaveKey, addr: u16, cnt: u16) -> Result<Vec<u16>> {
        let k = key.clone();
        self.ensure_ctx(&k).await?;
        let handle = self.get_handle(&k).await;
        let mut guard = handle.lock().await;
        let r = match guard.as_mut() {
            Some(ctx) => ctx.read_input_registers(addr, cnt).await.map_err(|e| anyhow!(e)),
            None => Err(anyhow!("No context")),
        };
        if r.is_err() {
            *guard = None;
        }
        r
    }

    pub async fn write_single(&self, key: SlaveKey, addr: u16, value: u16) -> Result<()> {
        let k = key.clone();
        self.ensure_ctx(&k).await?;
        let handle = self.get_handle(&k).await;
        let mut guard = handle.lock().await;
        let r = match guard.as_mut() {
            Some(ctx) => ctx.write_single_register(addr, value).await.map_err(|e| anyhow!(e)),
            None => Err(anyhow!("No context")),
        };
        if r.is_err() {
            *guard = None;
        }
        r.map(|_| ())
    }

    pub async fn write_multiple(&self, key: SlaveKey, addr: u16, values: &[u16]) -> Result<()> {
        let k = key.clone();
        self.ensure_ctx(&k).await?;
        let handle = self.get_handle(&k).await;
        let mut guard = handle.lock().await;
        let r = match guard.as_mut() {
            Some(ctx) => ctx.write_multiple_registers(addr, values).await.map_err(|e| anyhow!(e)),
            None => Err(anyhow!("No context")),
        };
        if r.is_err() {
            *guard = None;
        }
        r.map(|_| ())
    }

    pub async fn remove(&self, key: &SlaveKey) {
        let mut map = self.inner.write().await;
        if map.remove(key).is_some() {
            log::info!("🗑️  removed connection to {} (unit {})", key.addr, key.unit_id);
        }
    }

    pub async fn clear(&self) {
        let mut map = self.inner.write().await;
        map.clear();
        log::info!("🧹 cleared all modbus connections");
    }

    pub async fn connection_count(&self) -> usize {
        let map = self.inner.read().await;
        map.len()
    }
}

impl Default for ModbusConnectionPool {
    fn default() -> Self {
        Self::new()
    }
}